from pathlib import Path
from omegaconf import OmegaConf
from omegaconf.dictconfig import DictConfig
from numpy.random import seed as npseed
from random import seed as rseed
from torch import manual_seed, load
from monai.utils.misc import set_determinism
from omegaconf.listconfig import ListConfig
from omegaconf.dictconfig import DictConfig
from typing import Literal
from pytorch_lightning.loggers import WandbLogger
import json

import wandb.util

def get_training_type(cfg):
    """
    Get training type name from the config. 
    """
    training_type = cfg.general.get('training_type', 'autoencoder')
    if not isinstance(training_type, str):
        training_type = training_type.name
    return training_type

def init_experiment(cfg):
    print("#" * 40 + " Experiment: {} ".format(cfg.name) + "#" * 40)
    # seed_everything(cfg.general.seed)
    manual_seed(cfg.general.seed)
    npseed(cfg.general.seed)
    rseed(cfg.general.seed)
    set_determinism(seed=cfg.general.seed)

    if not hasattr(cfg.general, "run_id") or not cfg.general.run_id:
        cfg.general.run_id = wandb.util.generate_id()

    cfg.general.log_pth = cfg.general.log_pth.expanduser().resolve()    
    cfg.general.log_pth = cfg.general.log_pth / cfg.name / cfg.general.run_id
    cfg.general.log_pth.mkdir(parents=True, exist_ok=True)
   
    print("Log directory created: {}".format(cfg.general.log_pth))

    save_config(cfg)

    return cfg

def save_config(cfg):
    """
    Turn nested EasyDicts & PosixPaths into yaml readable format and save
    """
    cfg_dict = _recursive_confdict2dict(cfg)
    OmegaConf.save(DictConfig(cfg_dict), cfg.general.log_pth / "config.yaml")


def _recursive_confdict2dict(cfg_dict):
    cfg_dict_new = dict(cfg_dict.copy())
    for k, v in cfg_dict.items():
        if isinstance(v, DictConfig):
            cfg_dict_new[k] = _recursive_confdict2dict(v)
        elif isinstance(v, Path):
            cfg_dict_new[k] = str(v.expanduser().resolve())
        
    return cfg_dict_new

def wandblogger(cfg, log_path):
    with open('wandb_setup.json', 'r') as f:
        wandb_setup = json.load(f) 
    
    return WandbLogger(
        project=wandb_setup['project'],
        save_dir=log_path,
        name=cfg.name,
        log_model=False,
        entity=wandb_setup['entity'],
        config=dict(cfg),
        id=cfg.general.run_id,
    )

def resolve_listconfig_of_dicts(list_cfg: ListConfig):
    if any([isinstance(key_value, DictConfig) for key_value in list_cfg]) is False:
        print(f"resolve listconfig that is not a list of dict, return as is {list_cfg=}")
        return list_cfg
    return {
        key: value
        for key_value in list_cfg
        for key, value in key_value.items()
    }

def get_adapter_opts(cfg, task: str, adapter_type: Literal['input', 'output'] = 'output'):
    model_params_cfg = cfg.model.model_params
    adapter_cfg = model_params_cfg.get('output_tasks') if adapter_type == 'output' else model_params_cfg.get('input_tasks')
    adapter_cfg = resolve_listconfig_of_dicts(adapter_cfg)
    adapter_opts = adapter_cfg[task]
    if isinstance(adapter_opts, ListConfig):
        adapter_opts = resolve_listconfig_of_dicts(adapter_opts)
    if isinstance(adapter_opts, DictConfig):
        raise NotImplementedError
    if isinstance(adapter_opts, str):
        adapter_opts = {}
    adapter_opts = {k: v for k, v in adapter_opts.items() if k != 'type'} 
    return adapter_opts